#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#include <random>
#include <string>
#include <time.h>
#include <ctime>
#include <tuple>
#include <array>

/**
 * 超参数定义
 */

std::vector<int> real_coefficient();
const int Data_Size = 200;
const int Feature_Size = 4;
const double Test_Size = 0.3;
std::array<double, Feature_Size> Weights;
const std::vector<int> Coefs = real_coefficient();
const int Max_Iter = 3000;
std::vector<double> Cost;
const double Alpha = 0.01;

struct DataSet
{
    std::vector<std::vector<double>> x;
    std::vector<double> y;
};

int random_int(const int &min, const int &max)
{
    std::default_random_engine e;
    std::uniform_int_distribution<int> u(min, max); // 左闭右闭区间
    e.seed(time(0));
    return u(e);
}

/**
 * 生成数据集的特征的系数
 */
std::vector<int> real_coefficient()
{
    std::vector<int> Coefs;
    for (size_t i = 0; i < Feature_Size; i++)
    {
        Coefs.push_back(random_int(1, 10));
    }
    return Coefs;
}

double random_normal()
{
    // 生成随机数
    std::mt19937 rd;
    rd.seed(std::random_device()());
    // std::uniform_real_distribution<double> distribution(-1, 1);
    std::normal_distribution<double> normal_distribution(0, 1);
    return normal_distribution(rd);
}

std::tuple<DataSet, DataSet> create_daatset()
{
    DataSet ds_train, ds_test;
    // std::vector<int> Coefs = real_coefficient();
    int train_size = (int)Data_Size * Test_Size;

    for (size_t i = 0; i < Data_Size; i++)
    {
        std::vector<double> temp_x;
        double temp_y = 0;
        for (size_t j = 0; j < Feature_Size; j++)
        {
            double x_rd = random_normal();
            temp_x.push_back(x_rd);
            temp_y += x_rd * Coefs[j];
        }
        if (i < train_size)
        {
            ds_train.x.push_back(temp_x);
            ds_train.y.push_back(temp_y);
        }
        else
        {
            ds_test.x.push_back(temp_x);
            ds_test.y.push_back(temp_y);
        }
    }
    return std::make_tuple(ds_train, ds_test);
}

void print_result()
{
    std::cout << "------------------------------------" << std::endl;
    std::cout << "---------------- real Coefs --------------" << std::endl;
    for (size_t i = 0; i < Feature_Size; i++)
    {
        std::cout << Coefs[i] << " ";
    }
    std::cout << std::endl;
    std::cout << "---------------- fit Weights --------------" << std::endl;
    for (size_t i = 0; i < Feature_Size; i++)
    {
        std::cout << Weights[i] << " ";
    }
    std::cout << std::endl;
}

void print_cost(const std::vector<double> &cost)
{
    for (size_t i = 0; i < cost.size(); i++)
    {
        std::cout << i << " --> " << cost[i] << std::endl;
    }
}

double mse(const DataSet &dataset)
{
    std::vector<std::vector<double>> x = dataset.x;
    std::vector<double> y = dataset.y;
    double mse = 0;
    for (size_t i = 0; i < x.size(); i++)
    {
        double temp_y = 0;
        for (size_t j = 0; i < Feature_Size; i++)
        {
            temp_y += Weights[j] * x[i][j];
        }
        mse += pow(y[i] - temp_y, 2);
    }
    return pow(mse, 0.5);
}

void gradient_descent_step(const DataSet &dataset)
{
    // 初始化模型参数
    std::vector<double> sum_grad_w, grad_w;

    int m = dataset.x.size();
    for (size_t i = 0; i < Feature_Size; i++)
    {
        Weights[i] = random_normal();
        sum_grad_w.push_back(0);
        grad_w.push_back(0);
    }
    for (size_t i = 0; i < m; i++)
    {
        std::vector<double> x = dataset.x[i];
        double y = dataset.y[i];
        for (size_t j = 0; j < x.size(); j++)
        {
            sum_grad_w[j] += (Weights[j] * x[j] - y) * x[j];
        }
    }
    for (size_t i = 0; i < sum_grad_w.size(); i++)
    {
        grad_w[i] = 2 / m * sum_grad_w[i];
        Weights[i] -= Alpha * grad_w[i];
    }
}

void gradient_descent(const DataSet &dataset)
{
    std::vector<double> cost;
    for (size_t i = 0; i < Max_Iter; i++)
    {
        gradient_descent_step(dataset);
        cost.push_back(mse(dataset));
        print_result();
    }
    // print_cost(cost);
    // print_result();
}

int main()
{
    std::tuple<DataSet, DataSet> ds = create_daatset();
    DataSet train_ds = std::get<0>(ds);
    gradient_descent(train_ds);
    std::cout << "linear regression impl by gradient descend ..." << std::endl;
    return 0;
}